import sys
import csv
import time
import paramiko
import threading
import signal
from setting import *
import json
import requests
from datetime import datetime, timedelta
from use_case import cases, duration, start_dst_port

DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S"
URL = "http://10.112.39.63:8008/api/background/add"
STOP = False

channels = []
sshs = []
src_hosts = []
dst_hosts = []


def execute_command(hostname, port, username, password, command):
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    channel = None
    try:
        print(f'====================={hostname}======================')
        print(command)
        ssh.connect(hostname, port=port, username=username, password=password)
        channel = ssh.invoke_shell()
        channel.send(command)
    except paramiko.AuthenticationException:
        print(f"Failed to authenticate to {hostname}.")
    except paramiko.SSHException as e:
        print(f"SSH error occurred: {str(e)}")
    finally:
        return channel,ssh

def signal_handler(signal, frame):
    # 自定义信号处理函数
    print("接收到 Ctrl+C 信号，程序即将退出")
    global STOP
    STOP = True
signal.signal(signal.SIGINT, signal_handler)


def send_intent_request(src_ip,dst_ip,dst_port,service_type,slice_type,throughput_mbps):
    now = datetime.now()
    start_datetime = now.strftime(DATETIME_FORMAT)
    end_datetime = (now + timedelta(minutes=duration)).strftime(DATETIME_FORMAT)
    data=json.dumps({
        "src_ip":src_ip,
        "dst_ip":dst_ip,
        "dst_port":dst_port,
        "service_type":service_type,
        "slice_type":slice_type,
        "up_throughput":throughput_mbps,
        "down_throughput":0.01,
        "delay":10,
        "loss":0.1,
        "start_datetime":start_datetime,
        "end_datetime":end_datetime
        })
    headers = {
        "Content-Type": "application/json"
        }
    # 发送 POST 请求
    print(data)
    response = requests.post(URL, data=data, headers=headers)
    if response.status_code == 200:
        print("请求成功！")
    else:
        print("请求失败，错误码:", response.status_code)


def iperf(src, dst, dst_ip, dst_port, throughput):
    iperf_duration_second = (duration-1)*60

    iperf_command = f"iperf -u -c {dst_ip} -p {dst_port} -b {throughput}m -t {iperf_duration_second} -i 2 -V -l 1000\n"
    server_command = f"iperf -u -s -p {dst_port} -i 2 -w 512m -V\n"

    s_channel,s_ssh = execute_command(hosts_manage_ip[dst], ssh_port, host_username, host_password, server_command)
    c_channel,c_ssh = execute_command(hosts_manage_ip[src], ssh_port, host_username, host_password, iperf_command)
    return s_channel,s_ssh,c_channel,c_ssh


def iperf_server(dst, dst_port):
    server_command = f"iperf -u -s -p {dst_port} -i 2 -w 512m -V\n"
    s_channel,s_ssh = execute_command(hosts_manage_ip[dst], ssh_port, host_username, host_password, server_command)
    return s_channel,s_ssh

def iperf_client(src, dst_ip, dst_port, throughput):
    iperf_duration_second = (duration-1)*60
    iperf_command = f"iperf -u -c {dst_ip} -p {dst_port} -b {throughput}m -t {iperf_duration_second} -i 2 -V -l 1000\n"
    print(iperf_command)
    c_channel,c_ssh = execute_command(hosts_manage_ip[src], ssh_port, host_username, host_password, iperf_command)
    return c_channel, c_ssh


def iperf_servers():
    index = 1
    dst_port = start_dst_port
    for src, dst, service_type, slice_type,throughput in cases:
        print(f'==================Lunach iperf server of flow {index}=============================+')
        src, dst = sw2host[src], sw2host[dst]
        src_ip, dst_ip = f'{host_data_ip_addr_prefix}{src}', f'{host_data_ip_addr_prefix}{dst}'
        send_intent_request(src_ip, dst_ip, dst_port, service_type, slice_type, throughput)
        s_channel,s_ssh  = iperf_server(dst, dst_port)
        channels.append(s_channel)
        sshs.append(s_ssh)
        src_hosts.append(src)
        dst_hosts.append(dst)
        dst_port += 1
        index += 1
        if STOP:
            break


def iperf_clients():
    index = 1
    dst_port = start_dst_port
    for src,dst,slice_type,throughput in cases:
        print(f'==================Lunach iperf client of flow {index}=============================+')
        src, dst = sw2host[src], sw2host[dst]
        dst_ip = f'{host_data_ip_addr_prefix}{dst}'
        s_channel,s_ssh  = iperf_client(src, dst_ip, dst_port, throughput)
        channels.append(s_channel)
        sshs.append(s_ssh)
        src_hosts.append(src)
        dst_hosts.append(dst)
        dst_port += 1
        index += 1
        if STOP:
            break    

def slow_iperfs():
    global STOP
    dst_port = start_dst_port
    try:
        index = 1
        for src,dst,service_type,slice_type,throughput in cases:
            print(f'============================Flow {index}====================================')
            src, dst = sw2host[src], sw2host[dst]
            src_ip = f'{host_data_ip_addr_prefix}{src:x}'
            dst_ip = f'{host_data_ip_addr_prefix}{dst:x}'
            send_intent_request(src_ip,dst_ip,dst_port,service_type,slice_type,throughput)
            s_channel,s_ssh,c_channel,c_ssh = iperf(src,dst,dst_ip,dst_port,throughput)
            channels.append(c_channel)
            channels.append(s_channel)
            sshs.append(c_ssh)
            sshs.append(s_ssh)
            src_hosts.append(src)
            dst_hosts.append(dst)
            dst_port += 1
            index += 1
            if STOP:
                break
        print("All background has send.")
    except KeyboardInterrupt:
        print("Received CTRL+C, exiting.")
        STOP = True


def exit():
    global STOP
    if STOP:
        return
    print("======= Exiting... ======")    
    for src,dst, channel,ssh in zip(src_hosts, dst_hosts, channels, sshs):
        channel.send('\x03')
        channel.send('\x03')
        output = channel.recv(400).decode()
        while channel.recv_ready():
            last = output
            output = channel.recv(400).decode()
            
        print(f"========================|| H{src}->H{dst} ||===================================")
        output = last + output
        print(output)
        channel.close()
        ssh.close()
    STOP = True
    sys.exit(0)

if __name__ == "__main__":
    # 操作选项和对应命令
    options = [
        slow_iperfs,
        exit
    ]

    id = -1
    hint = ""
    for i, option in zip(range(len(options)),options):
        hint += '[%d] %s\n'%(i, option.__name__)
    try:
        while not STOP:
            new_id = int(input("请输入操作选项的编号:\n"+hint))
            if id == 0 and new_id == 0:
                print("不要重复iperf打流")
                continue
            id = new_id
            options[id]()
    except Exception as e:
        print(e)
        exit()
    finally:
        exit()
        
            
        

        


